import os
import random
import sys
from argparse import ArgumentParser
from datetime import datetime
from multiprocessing import Pool, cpu_count, Manager
from os import path, getenv, chdir, getcwd
from pathlib import Path
from random import shuffle
from shutil import rmtree
from time import sleep

import requests
from loguru import logger
from tqdm import tqdm


class Download:
	def __init__(self, txt):
		"""

		:param txt: URL清单文本文件路径
		"""
		self.txt = txt
		self.url_list = []
		self.url_dict = {}
		self.concurrency = getenv("Concurrency")
		self.domain_list = []
		self.rw = 1

	def init_read_txt(self):
		if self.txt is None:
			return
		f = Path(self.txt)
		logger.debug(f"正在判断文件: {self.txt}")
		if f.is_file():
			logger.debug(f"正在读取文件： {self.txt}")
			try:
				with open(file=self.txt, encoding="utf8", mode="r") as r:
					txt = r.read()
					ul = str(txt).split("\n")
					if len(ul) >= 1:
						s = 1
						for url in ul:
							if url not in self.url_list:
								self.url_list.append(s)
								self.url_dict[s] = url
								s += 1
			except Exception as e:
				logger.error(str(e))
		else:
			logger.error(f"找不到URL清单文件,请检查参数")
			sys.exit(3)
		logger.info(f"当前获取URL总数: {len(self.url_list)}")

	def init_config(self):
		"""
		初始化变量
		:return:
		"""
		self.url_list.append(
			"https://mirrors.aliyun.com/deepin-cd/23-Beta2/deepin-desktop-community-23-Beta2-amd64.iso")
		self.url_list.append(
			"https://mirrors.aliyun.com/centos/8-stream/isos/aarch64/CentOS-Stream-8-20231009.1-aarch64-boot.iso")
		self.url_list.append("https://mirrors.aliyun.com/ubuntu-releases/mantic/ubuntu-23.10.1-desktop-amd64.iso")
		self.url_list.append("https://mirrors.aliyun.com/archlinux/iso/2023.10.14/archlinux-2023.10.14-x86_64.iso")
		self.url_list.append(
			"https://mirrors.tuna.tsinghua.edu.cn/kali-images/current/kali-linux-2023.3-installer-amd64.iso")
		self.url_list.append(
			"https://mirrors.tuna.tsinghua.edu.cn/nixos-images/nixos-23.05/latest-nixos-plasma5-x86_64-linux.iso")

	def wget(self, url, index):
		"""
		开始下载
		:return:
		"""
		filename = f"file_{index}"
		logger.info(f"开始下载: [ {url} -> [ {filename} ]")
		host = str(url).split("//")[-1].split("/")[0]
		if host in self.domain_list:
			logger.warning(f"已存在相同域名[ {host} ]任务,本次任务取消...")
			return
		headers = {
			'Connection': 'keep-alive',
			'sec-ch-ua': '"Chromium";v="21", " Not;A Brand";v="99"',
			'sec-ch-ua-mobile': '?0',
			'sec-ch-ua-platform': '"Windows"',
			'Host': host,
			'Upgrade-Insecure-Requests': '1',
			'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/95.0.4638.69 Safari/537.36',
			'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9',
			'Sec-Fetch-Site': 'none',
			'Sec-Fetch-Mode': 'navigate',
			'Sec-Fetch-User': '?1',
			'Sec-Fetch-Dest': 'document',
			'Accept-Language': 'zh-CN,zh;q=0.9',
		}

		response = requests.get(url=url, headers=headers, stream=True, timeout=5)

		total_length = response.headers.get('content-length')
		if total_length is None:  # 如果服务器不返回内容长度，则无法显示进度
			try:
				with open(filename, 'wb') as f:
					for data in response.iter_content(chunk_size=8192):
						if data:
							f.write(data)
			except Exception as err:
				logger.error(str(err))
				logger.info(f"下载失败: [ {url} -> [ {filename} - {str(err)}]")
				return
		else:
			try:
				with open(filename, 'wb') as f, tqdm(total=int(total_length), desc='Downloading') as pbar:
					for data in response.iter_content(chunk_size=8192):
						if data:
							f.write(data)
							pbar.update(len(data))
			except Exception as err:
				logger.error(str(err))
				logger.info(f"下载失败: [ {url} -> [ {filename} - {str(err)}]")
				return
		logger.info(f"下载成功: [ {url} -> [ {filename} ]")
		if os.system(f"rm -f {filename}") != 0:
			logger.error(f"删除失败: {filename}")
		return

	def start(self):
		logger.debug(f"当前设置最小间隔时间: {arg_mini}")
		logger.debug(f"当前设置最大间隔时间: {arg_max}")
		logger.debug(f"当前设置URL文件路径: {t}")
		logger.debug(f"当前设置文件下载路径: {p}")
		self.init_read_txt()
		try:
			concurrency = int(self.concurrency)
		except Exception as err:
			logger.error(f"无法转换数据类型: {err}")
			sys.exit(2)
		logger.debug(f"当前设置任务数量: {concurrency}")

		self.rw = 1
		shuffle(self.url_list)
		single = False
		if str(getenv("Single")) == str('1'):
			single = True
		if int(concurrency) == 1 or int(concurrency) == 0 or single:
			index = random.randint(1, len(self.url_list))
			url = self.url_dict[index]
			filename = f"file_{index}"
			os.system(f"rm -f {filename}")
			logger.debug(f"正在执行第[ {self.rw} ]个任务")
			self.wget(url=url, index=index)
			os.system(f"rm -f {filename}")
		else:
			# 需要先在主函数定义局部变量
			pool = Pool(cpu_count())
			logger.debug("使用线程池模式")
			for ii in range(int(concurrency)):
				index = random.randint(1, len(self.url_list))
				url = self.url_dict[index]
				logger.debug(f"正在执行第[ {self.rw} ]个任务")
				pool.apply_async(func=self.wget, args=(url, index,))
				self.rw += 1
			pool.close()
			pool.join()


if __name__ == "__main__":
	v = getenv("Version")
	arg = ArgumentParser(description=f'当前脚本版本: {v}', prog="流量消失工具-Linux")
	# home = getenv("HOME")
	logger.info(f"当前脚本版本: {v}")
	home = getcwd()
	p = path.join(home, "tmp", "docker")
	time_max = getenv("Time_max")
	time_mini = getenv("Time_mini")
	concurrency_ = getenv("Concurrency")
	if time_max is None:
		time_max = 100
	if time_mini is None:
		time_mini = 50
	if concurrency_ is None:
		concurrency_ = 5
	h = f'[必选]设置URL清单文件(一行一个URL)'
	arg.add_argument('-t', '--txt', type=str, help=h, required=True)

	h = f'设置下载存放路径,默认: {p}'
	arg.add_argument('-p', '--path', type=str, default=p, help=h, required=False)
	h = '使用后台持续运行(程序根据设置的任务时间区间进行)'
	arg.add_argument('-d', action='store_true', help=h)

	h = "设置任务间隔时间随机最小数(单位: 秒),默认: 100"
	arg.add_argument('-mini', type=int, default=time_mini, help=h, required=False)
	h = "设置任务间隔时间随机最大数(单位: 秒),默认: 1000"
	arg.add_argument('-max', type=int, default=time_max, help=h, required=False)

	h = "设置多任务数量,默认: 5"
	arg.add_argument('-c', '--concurrency', type=int, default=concurrency_, help=h, required=False)

	h = "设置任务开始区间时间，默认: 7(早上7点开始)"
	arg.add_argument('-m', '--morning', type=int, default=7, help=h, required=False)
	h = "设置任务开始区间时间，默认: 23(晚上23点结束)"
	arg.add_argument('-n', '--night', type=int, default=23, help=h, required=False)

	h = '进程检测是否使用-ef参数(OpenWrt不需要该参数,主流Linux需要)'
	arg.add_argument('-ef', action='store_true', help=h)

	h = "判断时间单位增加,用于调整容器时间与实际实际不一致的问题,默认: +8"
	arg.add_argument('-a', '--add', type=int, default=8, help=h, required=False)

	args = arg.parse_args()
	p = args.path
	ef = args.ef
	t = args.txt
	arg_mini = args.mini
	arg_max = args.max
	arg_d = args.d
	arg_add = args.add
	arg_concurrency = args.concurrency
	# 时间参数

	arg_morning = args.morning
	arg_night = args.night
	now = datetime.now()
	hour = now.hour + int(arg_add)
	if arg_night > 24:
		logger.error(f"最晚时间设置异常: {arg_night}")
		sys.exit(2)
	if arg_morning < 0:
		logger.error(f"最早时间设置异常: {arg_morning}")
		sys.exit(2)
	if arg_morning >= arg_night:
		logger.error(f"当前设置时间异常,最早时间不能大于或等于最晚时间")
		sys.exit(2)
	if ef:
		logger.debug("启用进程检测为ef模式...")
	else:
		logger.debug("使用Openwrt运行模式...")
	if arg_d:
		logger.debug("当前使用后台持续运行(容器推荐)模式...")
	else:
		logger.debug("当前使用前台运行(调试推荐)模式...")
	if path.isdir(p):
		chdir(p)
		try:
			for i in os.listdir("."):
				i = path.abspath(i)
				if path.isdir(i):
					logger.debug(f"正在删除: {i}")
					rmtree(path=i)
				else:
					os.remove(i)
		except Exception as e:
			logger.error(e)
			sys.exit(3)
		if arg_d:
			while True:
				d = Download(txt=t)
				d.concurrency = arg_concurrency
				if int(arg_morning) <= int(hour) <= int(arg_night):
					logger.debug(f"当前时间[ {hour} ]符合设置区间[ {arg_morning} -> {arg_night} ],任务执行...")
					d.start()
					sleep(random.randint(arg_mini, arg_max))
				else:
					sec = int(arg_night - arg_morning) * 60
					mess = f"[ {v} ]当前时间[ {hour} ]不在设定任务时间内[ {arg_morning} -> {arg_night} ],等待{sec}秒后继续..."
					logger.debug(mess)
					sleep(sec)
		else:
			if int(arg_morning) <= int(hour) <= int(arg_night):
				logger.debug(f"当前时间[ {hour} ]符合设置区间[ {arg_morning} -> {arg_night} ],任务执行...")
				d = Download(txt=t)
				d.concurrency = arg_concurrency
				d.start()
			else:
				logger.debug(f"[ {v} ]当前时间[ {hour} ]不在设定任务时间内[ {arg_morning} -> {arg_night} ]...")
	else:
		logger.error(f"[ {v} ]路径不存在: {p}")
		sys.exit(4)
